import os
os.environ['HF_ENDPOINT']='https://hf-mirror.com'
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel
import os
import json
from tqdm import tqdm  # 用于显示美观的进度条

# --- 1. 配置部分 (请根据您的实际情况修改) ---

# 您的预训练基础模型名称 (必须与训练时使用的模型一致)
BASE_MODEL_NAME = "hfl/chinese-roberta-wwm-ext"

# 您想要加载的模型checkpoint文件夹的路径
CHECKPOINT_PATH = "./results_sampled/checkpoint-10016"  # <--- 请修改

# 存放源JSON文件的文件夹路径
INPUT_FOLDER = r"C:\Users\Administrator\Projects\lmy\outputs\step2_sentence_extraction\t"  # <--- 请修改

# 结果要保存到的目标JSON文件路径
OUTPUT_JSON_PATH = "../outputs/step2_sentence_extraction/t2/filtered_results_list.json"  # <--- 请修改

# 每个类别需要收集的目标句子数量
TARGET_COUNT = 5000

# 批量预测时，每个批次处理的句子数量 (可根据您的显存大小调整)
BATCH_SIZE = 64

# --- 2. 初始化模型和分词器 ---

# 判断使用GPU还是CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"正在使用设备: {device}")

# 加载分词器
print("正在加载分词器...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)

# 加载基础模型
print("正在加载基础模型...")
base_model = AutoModelForSequenceClassification.from_pretrained(
    BASE_MODEL_NAME,
    num_labels=3
)

# 加载并融合LoRA权重
print(f"正在从 '{CHECKPOINT_PATH}' 加载LoRA适配器...")
model = PeftModel.from_pretrained(base_model, CHECKPOINT_PATH)
model.to(device)
model.eval()  # 设置为评估模式
print("模型加载完成，已准备就绪。")


# --- 3. 定义批量预测函数 ---

def predict_batch(sentences: list):
    """
    对一个批次的句子进行分类预测。
    """
    if not sentences:
        return []

    with torch.no_grad():
        # 使用分词器处理整个批次
        inputs = tokenizer(
            sentences,
            padding=True,  # 填充到批次中最长句子的长度
            truncation=True,  # 截断超过模型最大长度的句子
            max_length=512,
            return_tensors="pt"  # 返回PyTorch张量
        )
        inputs = {key: value.to(device) for key, value in inputs.items()}
        outputs = model(**inputs)
        predictions = torch.argmax(outputs.logits, dim=1)
        # 将结果从GPU转移到CPU，并转换为列表
        return predictions.cpu().tolist()


# --- 4. 主处理逻辑 ---

print("\n--- 开始处理文件 ---")

# 准备存储结果的列表
results_class_0 = []
results_class_1 = []

# 获取所有源JSON文件的路径
try:
    json_files = [os.path.join(INPUT_FOLDER, f) for f in os.listdir(INPUT_FOLDER) if f.endswith('.json')]
    if not json_files:
        print(f"错误：在文件夹 '{INPUT_FOLDER}' 中没有找到任何.json文件。")
        exit()
except FileNotFoundError:
    print(f"错误：找不到文件夹 '{INPUT_FOLDER}'。请检查路径是否正确。")
    exit()

# 主循环，用于遍历所有文件，直到收集满所需数量的句子
try:
    stop_processing = False
    for file_path in json_files:
        if stop_processing:
            break

        print(f"\n正在处理文件: {file_path}")

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                sentences = json.load(f)
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            print(f"  - 文件 {file_path} 读取或解析失败，已跳过。错误: {e}")
            continue

        if not isinstance(sentences, list):
            print(f"  - 文件 {file_path} 格式不正确（不是一个JSON列表），已跳过。")
            continue

        # 使用tqdm创建进度条
        progress_bar = tqdm(range(0, len(sentences), BATCH_SIZE), desc="  - 预测进度")

        for i in progress_bar:
            if stop_processing:
                break

            # 获取当前批次的句子
            batch_sentences = sentences[i:i + BATCH_SIZE]

            # 进行批量预测
            batch_predictions = predict_batch(batch_sentences)

            # 遍历批次结果
            for sentence, pred_class in zip(batch_sentences, batch_predictions):
                if pred_class == 0 and len(results_class_0) < TARGET_COUNT:
                    results_class_0.append(sentence)
                elif pred_class == 1 and len(results_class_1) < TARGET_COUNT:
                    results_class_1.append(sentence)

            # 更新进度条的后缀信息
            progress_bar.set_postfix({
                'Class 0': f'{len(results_class_0)}/{TARGET_COUNT}',
                'Class 1': f'{len(results_class_1)}/{TARGET_COUNT}'
            })

            # 检查是否已达到目标数量
            if len(results_class_0) >= TARGET_COUNT and len(results_class_1) >= TARGET_COUNT:
                print("\n已成功收集到足够数量的句子，处理提前结束。")
                stop_processing = True
                break

except KeyboardInterrupt:
    print("\n用户手动中断处理。")
except Exception as e:
    print(f"\n处理过程中发生未知错误: {e}")

# --- 5. 保存结果 (已修改为单一列表格式) ---

print("\n--- 处理完成，正在保存结果 ---")

# *** 修改部分 ***
# 将两个列表合并成一个单一的大列表
final_results_list = results_class_0 + results_class_1

# 保存到JSON文件
try:
    with open(OUTPUT_JSON_PATH, 'w', encoding='utf-8') as f:
        # 直接将这个列表写入JSON文件
        json.dump(final_results_list, f, ensure_ascii=False, indent=4)

    print(f"结果已成功保存到: {OUTPUT_JSON_PATH}")
    print(f"  - 类别0句子数量: {len(results_class_0)}")
    print(f"  - 类别1句子数量: {len(results_class_1)}")
    print(f"  - 文件中总句子数量: {len(final_results_list)}")
except Exception as e:
    print(f"保存文件时出错: {e}")